"""
物理感知损失函数
功能：将材料力学定律作为正则化项约束模型
"""
import torch


class PhysicsLoss:
    def __init__(self, concrete_strength={'C30': 30, 'C40': 32.5, 'C50': 34.5}):
        """
        参数说明：
        - concrete_strength: 混凝土弹性模量（GPa→转换为MPa）
        """
        self.E_map = {k: v * 1e3 for k, v in concrete_strength.items()}  # GPa→MPa

    def __call__(self, pred_damage, struct_params, strain):
        """
        计算物理约束损失
        参数：
        - pred_damage: 模型预测损伤指数 [batch]
        - struct_params: 结构参数 [batch, features]
        - strain: 应变测量值 [batch]
        """
        # 获取混凝土弹性模量
        E = torch.zeros_like(pred_damage)
        for i, grade_idx in enumerate(torch.argmax(struct_params[:, 3:6], dim=1)):
            # 假设结构参数最后3位是混凝土独热编码
            grade = list(self.E_map.keys())[grade_idx]
            E[i] = self.E_map[grade]

        # 理论应力：σ = E * ε
        stress_true = E * strain

        # 预测有效应力（损伤导致刚度下降）
        stress_pred = stress_true * (1 - pred_damage)

        # 损失项：应力差异的MSE
        return torch.mean((stress_pred - stress_true) ** 2)